#include "..\d3d.h"
#include "d3d_shader.h"
#include "..\d3d_texture.h"
#include "..\d3d_mesh.h"
#include "..\..\common\string.h"
#include "..\..\common\def.h"
#include "..\..\common\util.h"
#include "..\..\engine\engine.h"
#include "..\interface_types.h"
#include "..\..\engine\memorymanager.h"
#include "..\..\engine\error.h"
#include "..\..\math\math.h"


#include "..\..\common\types.h"
#include "..\..\common\def.h"
#include "..\..\common\arraylist.hpp"

#include "d3d_watershader.h"
D3DVERTEXELEMENT9 InstanceDecl[] =
{
	{0,0,D3DDECLTYPE_FLOAT3, D3DDECLMETHOD_DEFAULT, D3DDECLUSAGE_POSITION,0},
	{0,sizeof(float) * 3, D3DDECLTYPE_FLOAT3, D3DDECLMETHOD_DEFAULT, D3DDECLUSAGE_NORMAL,0},
	{0,sizeof(float) * 6,D3DDECLTYPE_FLOAT3, D3DDECLMETHOD_DEFAULT, D3DDECLUSAGE_TEXCOORD,0},
	{0,sizeof(float) * 9,D3DDECLTYPE_FLOAT2, D3DDECLMETHOD_DEFAULT, D3DDECLUSAGE_TEXCOORD,1},
	{0,sizeof(float) * 11,D3DDECLTYPE_FLOAT2, D3DDECLMETHOD_DEFAULT, D3DDECLUSAGE_TEXCOORD,2},
	
	{1,0,D3DDECLTYPE_FLOAT4, D3DDECLMETHOD_DEFAULT,D3DDECLUSAGE_TEXCOORD,3},
	{1,16,D3DDECLTYPE_FLOAT4, D3DDECLMETHOD_DEFAULT,D3DDECLUSAGE_TEXCOORD,4},
	{1,32,D3DDECLTYPE_FLOAT4, D3DDECLMETHOD_DEFAULT,D3DDECLUSAGE_TEXCOORD,5},
	{1,48,D3DDECLTYPE_FLOAT4, D3DDECLMETHOD_DEFAULT,D3DDECLUSAGE_TEXCOORD,6},
	D3DDECL_END()
}; 

LPD3DXEFFECTPOOL D3DShader::m_pEffectPool = NULL;
FogData D3DShader::m_Fog = FogData();

/**************************************
 *
 **************************************/
HRESULT D3DShader::CreateEffectPool()
{
	if(!ISNULL(D3DShader::m_pEffectPool))
		ReleaseEffectPool();
	
	if(FAILED(D3DXCreateEffectPool(&D3DShader::m_pEffectPool)))
	{
		EE_LOGERROR("Failed to Create effect pool.");
		return E_FAIL;
	}
	
	return S_OK;
}

/**************************************
 *
 **************************************/
LPD3DXEFFECTPOOL D3DShader::GetEffectPool()
{
	return D3DShader::m_pEffectPool;
}

/**************************************
 *
 **************************************/
HRESULT D3DShader::ReleaseEffectPool()
{
	if(ISNULL(D3DShader::m_pEffectPool))
	{
		EE_LOGERROR("No effect pool to Release.");
		return E_FAIL;
	}
	
	D3DShader::m_pEffectPool->Release();
	D3DShader::m_pEffectPool = NULL;
	return S_OK;
}

/**************************************
 * Used when loading shaders, this orders
 * the shader list so that it can be itterated
 * elegantly during runtime while producing
 * correct visual results
 **************************************/
bool D3DShader::funcRenderOrderSort(D3DShader* pA, D3DShader* pB)
{
	return (pA->RenderOrder() > pB->RenderOrder());
}

/**************************************
 *
 **************************************/
D3DShader::D3DShader(CD3D9Device* const pDevice) : IResource(pDevice)
{	
	m_pVertexDeclaration = NULL;
	m_pProjectionHandle = NULL; 
	m_pFogHandle = NULL;
	m_pViewHandle = NULL; 
	m_pTimeHandle = NULL; 
	m_pInstanceVB = NULL;
	m_pMeshVB = NULL;
	m_pEffect = NULL; 
	m_Identity = "D3DShader (unnamed)"; 
	m_lastInstanceCount = 0; 
	m_SubmittedVertexCount = 0;
	m_MeshCount = 0; 
	m_TotalInstanceCount = 0; 
	m_vbInstanceSize = 0;
}

/**************************************
 *
 **************************************/
D3DShader::~D3DShader()
{	
	Release();
}

/**************************************
 *
 **************************************/
string D3DShader::GetType() const
{
	return string("D3DShader"); 
}

/**************************************
 *
 **************************************/
unsigned int D3DShader::RenderOrder()
{
	return 0;
}

/**************************************
 *
 **************************************/
D3DVERTEXELEMENT9* D3DShader::GetVertexDeclarationElements()
{
	return InstanceDecl;
}

/**************************************
 * Initalizes the resource
 **************************************/
HRESULT D3DShader::Create(const string& szFilename)
{
	LPD3DXBUFFER pErrors	= NULL; 
	DWORD dwShaderFlags		= D3DXFX_NOT_CLONEABLE;
#ifdef _DEBUG
	dwShaderFlags |= D3DXSHADER_DEBUG; 
#endif; 

	try {

		if(ISNULL(GetDevice()))
			throw EE_EXCEPTION("Shader has no graphics device."); 

		if(!szFilename.IsValid())
			throw EE_EXCEPTION("Bad shader filename."); 


		if(ISNULL(m_pEffect)) {
			
			ReleaseResources();

			const string filepath = CEngine::GetInstance()->FindFilepath(szFilename); 
			if(!filepath.IsValid())
				throw EE_EXCEPTION("Failed to locate file."); 

			//Load the shader
			if(D3DXCreateEffectFromFile(GetDevice()->m_pD3DDevice,filepath.c_str(),NULL,NULL,
				dwShaderFlags,GetEffectPool(),&m_pEffect,&pErrors) != D3D_OK) {
					MessageBox(0, (char*)pErrors->GetBufferPointer(),0,0);
					throw EE_EXCEPTION(pErrors ? (char*)pErrors->GetBufferPointer() : "Unknown shader error"); 

			}
		}

		if(FAILED(InitalizeParameters(m_pEffect)))
			throw EE_EXCEPTION("Failed to initalize effect parameters."); 
		
		if(!ISNULL(m_pVertexDeclaration))
		{
			m_pVertexDeclaration->Release();
			m_pVertexDeclaration = NULL;
		}
		if(FAILED(GetDevice()->GetDevice()->CreateVertexDeclaration(GetVertexDeclarationElements(), &m_pVertexDeclaration)))
			throw EE_EXCEPTION("CreateVertexDeclaration(..) failed."); 

		m_Identity = szFilename; 
		return S_OK; 

	}catch(CError::CErrorData errorData)
	{
		CError::CError(errorData, ERRORREPORT_LOG); 
		Release(); 
		return E_FAIL; 
	}
}

/**************************************
 * Allows shader implementations to
 * specfiy custom parameters for an effect
 **************************************/
HRESULT D3DShader::InitalizeParameters(LPD3DXEFFECT pEffect)
{
	HRESULT hr;
	do {

		if(ISNULL(pEffect)) {
			EE_LOGERROR("Invalid parameter.");
			hr = E_INVALIDARG;
			break;
		}

		m_pProjectionHandle = pEffect->GetParameterBySemantic( NULL, "PROJECTION" ); 
		m_pViewHandle = pEffect->GetParameterBySemantic( NULL, "VIEW" ); 		
		m_pTimeHandle = pEffect->GetParameterBySemantic( NULL, "TIME" ); 
		m_pFogHandle = pEffect->GetParameterBySemantic( NULL, "FOG_DATA" ); 
		
		if(ISNULL(m_pProjectionHandle) || ISNULL(m_pViewHandle) || 
			ISNULL(m_pFogHandle) || ISNULL(m_pTimeHandle)) {
			EE_LOGERROR("Unknown semantic.");
			hr = E_UNEXPECTED;
			break;
		}
		return S_OK;
	
	}while(false);
	
	m_pProjectionHandle = NULL;
	m_pViewHandle = NULL;
	m_pTimeHandle = NULL;
	m_pFogHandle = NULL;
	return hr;
}

/**************************************
 * Updates the parameters of the shader
 **************************************/
HRESULT D3DShader::UpdateParameters(LPD3DXEFFECT pEffect)
{
	if(ISNULL(pEffect))
		return E_FAIL;

	ViewOrrientation* pCamera = &GetDevice()->m_camera; 

	pEffect->SetMatrix(m_pProjectionHandle,&pCamera->m_Projection); 
	pEffect->SetMatrix(m_pViewHandle,&pCamera->m_View); 
	pEffect->SetFloat(m_pTimeHandle,CEngine::GetInstance()->GetTimeDuration()); 
	return S_OK;
}

/**************************************
 * Allows client to submit geometry
 * that will have this shader be applied
 * to it during rendering
 **************************************/
void D3DShader::AddRenderQue(D3DRenderable* pMesh, const Matrix4f& transform)
{
	InstanceData instance; 
	instance.fData[0][0] = transform._M11; 
	instance.fData[0][1] = transform._M12; 
	instance.fData[0][2] = transform._M13; 
	instance.fData[0][3] = 0.0f; 

	instance.fData[1][0] = transform._M21; 
	instance.fData[1][1] = transform._M22; 
	instance.fData[1][2] = transform._M23; 
	instance.fData[1][3] = 0.0f; 

	instance.fData[2][0] = transform._M31; 
	instance.fData[2][1] = transform._M32; 
	instance.fData[2][2] = transform._M33; 
	instance.fData[2][3] = 0.0f; 

	instance.fData[3][0] = transform._M41; 
	instance.fData[3][1] = transform._M42; 
	instance.fData[3][2] = transform._M43; 
	instance.fData[3][3] = 1.0f; 
	
	m_SubmittedVertexCount += pMesh->getVertexCount();
	m_MeshHashInstances.Add(pMesh, instance); 
}

/**************************************
 * Constructs an instance buffer for this
 * shader based on the passed mesh instances
 * in the last frame
 **************************************/
HRESULT D3DShader::ConstructInstanceBuffer()
{
	unsigned int i = 0;
	unsigned int meshCount = 0; 
	unsigned int totalInstanceCount = 0; 
	m_MeshHashInstances.GetDataInfo(&meshCount, &totalInstanceCount); 

	if(meshCount == 0)
		return E_FAIL; 

	LPDIRECT3DDEVICE9 const pDevice = GetDevice()->m_pD3DDevice; 


	unsigned int mvbSize = m_SubmittedVertexCount * VertexComplex::size;
	
	if(m_pMeshVB) 
	{
		m_pMeshVB->Release(); 
		m_pMeshVB = NULL; 
	}
	//Build up VB of all meshes (they share the same vertex declaration)
	if(pDevice->CreateVertexBuffer(mvbSize,0,0, D3DPOOL_DEFAULT,&m_pMeshVB,NULL) != D3D_OK)
	{	
		m_vbInstanceSize = 0;
		m_MeshCount = 0;
		m_TotalInstanceCount = 0;
		return E_FAIL; 
	}

	//VertexComplex *pTmpMeshVertex = NULL; 
	VertexComplex *pTmpMeshVertex = NULL; 
	if(m_pMeshVB->Lock(0, NULL, (void **) &pTmpMeshVertex, 0) == D3D_OK)
	{
		for(unsigned int i = 0;  i < meshCount;  i++)
		{
			D3DRenderable*	pMesh = m_MeshHashInstances.KeyAtTableIndex(i);
			unsigned int	meshVertexCount = pMesh->getVertexCount();
			memcpy(pTmpMeshVertex, pMesh->getVertices(), meshVertexCount * pMesh->m_ElementSize); 
			pTmpMeshVertex += meshVertexCount; 
		}
	
		m_pMeshVB->Unlock(); 
	}
	else
	{	
		return E_FAIL; 
	}

	const unsigned int vbInstanceSize = sizeof(InstanceData); 
	const unsigned int vbSize = totalInstanceCount * vbInstanceSize; 
	
	if(totalInstanceCount > m_lastInstanceCount || !m_pInstanceVB)
	{
		m_lastInstanceCount = totalInstanceCount; 

		if(m_pInstanceVB) 
		{
			m_pInstanceVB->Release(); 
			m_pInstanceVB = NULL; 
		}

		if(pDevice->CreateVertexBuffer(vbSize,0,0, D3DPOOL_DEFAULT,&m_pInstanceVB,NULL) != D3D_OK)
		{	
			m_vbInstanceSize = 0;
			m_MeshCount = 0;
			m_TotalInstanceCount = 0;
			return E_FAIL; 
		}
	}

	InstanceData *pTmpInstTrans = NULL; 
	if(m_pInstanceVB->Lock(0, NULL, (void **) &pTmpInstTrans, 0) == D3D_OK)
	{
		InstanceData* pMeshBaseByte = pTmpInstTrans; 
		for(unsigned int i = 0;  i < meshCount;  i++)
		{
			HashTable<InstanceData, D3DRenderable*>::Pair *pPair = m_MeshHashInstances.GetPair(i); 		
			memcpy(pMeshBaseByte, pPair->m_pData, pPair->m_Count * vbInstanceSize); 
			pMeshBaseByte += pPair->m_Count; 
		}
		
		m_pInstanceVB->Unlock(); 
	}
	else
	{	
		m_vbInstanceSize = 0;
		m_MeshCount = 0;
		m_TotalInstanceCount = 0;
		return E_FAIL; 
	}

	m_vbInstanceSize = vbInstanceSize;
	m_MeshCount = meshCount;
	m_TotalInstanceCount = totalInstanceCount;
	return S_OK; 
} 

FogData D3DShader::GetFogParameters()
{
	return m_Fog;
}

void D3DShader::SetFogParameters(const FogData& fog)
{
	m_Fog = fog;
	m_pEffect->SetValue(m_pFogHandle,(void*)&m_Fog, sizeof(FogData)); 
}

void D3DShader::FlushRenderList(void)
{
	m_MeshHashInstances.ClearAllocatedSpace(); 
}

void D3DShader::SetTextures(D3DRenderable * const pMesh)
{
//	LPDIRECT3DDEVICE9 pRenderDevice = GetDevice()->GetDevice();

	CD3D9Device* pDevice = GetDevice();
	Texture* const pMeshTextureA = pMesh->getTexture(0);//m_pTextureUnitA;
	Texture* const pMeshTextureB = pMesh->getTexture(1);//m_pTextureUnitB;
	Texture* const pMeshTextureC = pMesh->getTexture(3);//m_pTextureMask;
	Texture* const pMeshTextureGI = pMesh->getTexture(4);//m_pTextureMask;

	pDevice->SetTexture(0,pMeshTextureA);
	pDevice->SetTexture(1,pMeshTextureB);
	pDevice->SetTexture(2,pMeshTextureC);
	pDevice->SetTexture(4,pMeshTextureGI);
	pDevice->SetTexture(5,pMesh->getTexture(2));
}

/**************************************
 * Appreciates that shaders can be executed
 * many times perframe. This provides opporutnity
 * to preform one-off initilization on a per-frame
 * basis
 **************************************/
void D3DShader::PreExecute()
{
	if(FAILED(UpdateParameters(m_pEffect)))
		return;
	//TODO: WORK OUT WHY THIS CANT BE DONE HERE
	//if(FAILED(ConstructInstanceBuffer(&m_vbInstanceSize, &m_MeshCount, &m_TotalInstanceCount)))
	//	return;
}

/**************************************
 * Provides oppourtunity for shader to 
 * execute code after shading, each frame
 **************************************/
void D3DShader::PostExecute()
{
	m_SubmittedVertexCount = 0;
	m_MeshHashInstances.ClearAllocatedSpace();
}

/**************************************
 * Executes the shader which usually results
 * in geometry being drawn
 **************************************/
void D3DShader::Execute(const int flags)
{
	if(!(flags & (RENDER_LIGHTING | RENDER_OPAQUE | RENDER_REFLECTED | RENDER_REFRACTED)))
		return; 

	unsigned int passes = 0;
	LPDIRECT3DDEVICE9 pDevice = GetDevice()->m_pD3DDevice; 
	
	if(FAILED(ConstructInstanceBuffer()))
		return;


	GetDevice()->SetCurrentVertexDeclaration(m_pVertexDeclaration);
	
	unsigned int instanceBytesOffset = 0; 
	unsigned int vertexOffset = 0;	
	pDevice->SetStreamSource(0, m_pMeshVB, 0, VertexComplex::size); 

	for(unsigned int i = 0;  i < m_MeshCount;  i++)
	{
		HashTable<InstanceData, D3DRenderable*>::Pair *pPair = m_MeshHashInstances.GetPair(i); 
		D3DRenderable* pMesh = pPair->m_Key;
		
		
		if(!(pMesh->getRenderFlags() & flags))
		{
			instanceBytesOffset += (pPair->m_Count * m_vbInstanceSize); 
			vertexOffset += pMesh->getVertexCount();
			continue; 
		}

		SetTextures(pMesh);
	
		pDevice->SetStreamSourceFreq(0, D3DSTREAMSOURCE_INDEXEDDATA | pPair->m_Count); 
	
		pDevice->SetStreamSource(1, m_pInstanceVB, instanceBytesOffset, m_vbInstanceSize); 
		pDevice->SetStreamSourceFreq(1, D3DSTREAMSOURCE_INSTANCEDATA | 1); 
		pDevice->SetIndices(pMesh->pIndexBuffer); 

		m_pEffect->Begin(&passes, 0); 
		for(unsigned int p = 0;  p < passes;  p++)
		{
			m_pEffect->BeginPass(p); 
			pMesh->Render(vertexOffset);
			m_pEffect->EndPass(); 
		}
		m_pEffect->End(); 
		instanceBytesOffset += (pPair->m_Count * m_vbInstanceSize); 
		vertexOffset += pMesh->getVertexCount();
	}

	pDevice->SetStreamSourceFreq(0, 1); 
	pDevice->SetStreamSourceFreq(1, 1); 
}

/**************************************
 * Enables an effect to be reloaded from
 * file during runtime
 **************************************/
HRESULT	D3DShader::Reload(void)
{
	LPD3DXBUFFER pErrors	= NULL; 
	LPD3DXEFFECT pEffect	= NULL;
	DWORD dwShaderFlags		= D3DXFX_NOT_CLONEABLE; 
#ifdef _DEBUG
	dwShaderFlags |= D3DXSHADER_DEBUG; 
#endif

	try {
	
		if(!m_Identity.IsValid() || ISNULL(m_pEffect))
			throw EE_EXCEPTION("Bad shader state.")

		const string filepath = CEngine::GetInstance()->FindFilepath(m_Identity); 
		if(!filepath.IsValid())
			throw EE_EXCEPTION("Failed to locate file."); 

		//Load the shader
		if(D3DXCreateEffectFromFile(GetDevice()->GetDevice(),filepath.c_str(),NULL,NULL,
			dwShaderFlags,GetEffectPool(),&pEffect,&pErrors) != D3D_OK)
			throw EE_EXCEPTION(pErrors ? (char*)pErrors->GetBufferPointer() : "Unknown shader error"); 

		if(FAILED(InitalizeParameters(pEffect)))
		{
			//Fix up parameters for origonal effect
			InitalizeParameters(m_pEffect);
			throw EE_EXCEPTION("Failed to initialize effect parameters during reload.")
		}
		
		//Seems safe to Release origonal effect
		ReleaseResources();
		m_pEffect = pEffect;

		return S_OK; 
	}catch(CError::CErrorData error)
	{
		CError::CError(error);
		if(!ISNULL(pEffect))
		{
			pEffect->Release();
			pEffect = NULL;
		}
		return E_FAIL; 
	}
}

/**************************************
 *Attempts to Release resoruces associated
 *with this shader instance
 **************************************/
void D3DShader::ReleaseResources()
{
	if(m_pEffect)
	{
		m_pEffect->Release(); 
		m_pEffect = NULL; 
	}

	if(m_pInstanceVB)
	{
		m_pInstanceVB->Release(); 
		m_pInstanceVB = NULL; 
	}
}

/**************************************
 *Attempts to Release resoruces associated
 *with this shader instance as well as itself
 *from the parent device
 **************************************/
HRESULT D3DShader::Release()
{
	m_IsValid = false;
	if(FAILED(GetDevice()->DeregisterResource(this)))
	{
		m_IsValid = true;
		return E_FAIL;
	}

	ReleaseResources();
	
	if(!ISNULL(m_pVertexDeclaration))
	{
		m_pVertexDeclaration->Release();
		m_pVertexDeclaration = NULL;
	}
	
	if(!ISNULL(m_pInstanceVB))
	{
		m_pInstanceVB->Release();
		m_pInstanceVB = NULL;
	}
	if(!ISNULL(m_pMeshVB))
	{
		m_pMeshVB->Release();
		m_pMeshVB = NULL;
	}

	return S_OK; 
}